import json
import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE

# 加载生成的嵌入向量
def load_embeddings(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    
    image_embeddings = []
    text_embeddings = []
    labels = []
    
    for item in data['train']:
        image_embeddings.append(item['image_embedding'])
        text_embeddings.append(item['report_embedding'])
        labels.append(item['id'])
    
    image_embeddings = np.array(image_embeddings)
    text_embeddings = np.array(text_embeddings)
    labels = np.array(labels)
    
    return image_embeddings, text_embeddings, labels

# 使用 t-SNE 进行降维
def perform_tsne(embeddings, n_components=2):
    tsne = TSNE(n_components=n_components, random_state=42)
    embeddings_tsne = tsne.fit_transform(embeddings)
    return embeddings_tsne

# 可视化降维后的数据
def visualize_embeddings(embeddings_tsne, labels, title):
    plt.figure(figsize=(10, 8))
    plt.scatter(embeddings_tsne[:, 0], embeddings_tsne[:, 1], c=labels, cmap='viridis', s=50)
    plt.colorbar()
    plt.title(title)
    plt.xlabel('t-SNE Component 1')
    plt.ylabel('t-SNE Component 2')
    plt.show()

# 主函数
def main():
    # 加载嵌入向量
    file_path = 'data/iu_xray/iu_xray/embeddings.json'
    image_embeddings, text_embeddings, labels = load_embeddings(file_path)

    # 对图像嵌入向量进行 t-SNE 降维
    image_embeddings_tsne = perform_tsne(image_embeddings)
    visualize_embeddings(image_embeddings_tsne, labels, 'Image Embeddings t-SNE Visualization')

    # 对文本嵌入向量进行 t-SNE 降维
    text_embeddings_tsne = perform_tsne(text_embeddings)
    visualize_embeddings(text_embeddings_tsne, labels, 'Text Embeddings t-SNE Visualization')

if __name__ == "__main__":
    main()